class LightningModel(pl.LightningModule):
def __init__(self, model, optimizer_name, optimizer_hparams):
"""LightningModel.
Args:
model: Name of the model/CNN to run. Used for creating the model
optimizer_name: Name of the optimizer to use. Currently supported: Adam, SGD
optimizer_hparams: Hyperparameters for the optimizer, as dictionary. This includes learning rate, weight decay, etc.
"""
super().__init__()
self.model = model
self.save_hyperparameters(ignore=['model'])
self.loss_module = torch.nn.MSELoss(reduction = 'mean')
# Set up attributes for computing the accuracy
self.train_psnr = torchmetrics.image.PeakSignalNoiseRatio()
self.train_ssim = torchmetrics.image.StructuralSimilarityIndexMeasure()
self.valid_psnr = torchmetrics.image.PeakSignalNoiseRatio()
self.valid_ssim = torchmetrics.image.StructuralSimilarityIndexMeasure()
self.test_psnr = torchmetrics.image.PeakSignalNoiseRatio()
self.test_ssim = torchmetrics.image.StructuralSimilarityIndexMeasure()
def forward(self, x):
return self.model(x)
# A common forward step to compute the loss and label
# this is used for training, validation, and testing below
def _shared_step(self, batch):
image, label = batch
predicted_label = self(image)
loss = self.loss_module(predicted_label, label)
return loss, label, predicted_label
def training_step(self, batch, batch_idx):
loss, label, predicted_label = self._shared_step(batch)
self.train_psnr.update(predicted_label, label)
self.train_ssim.update(predicted_label, label)
metrics = {'train_loss': loss, 'train_psnr': self.train_psnr, 'train_ssim': self.train_ssim}
self.log_dict(metrics, on_epoch=True, on_step=False, prog_bar=True)
return loss # this is passed to the optimzer for training
def validation_step(self, batch, batch_idx):
loss, label, predicted_label = self._shared_step(batch)
self.valid_psnr.update(predicted_label, label)
self.valid_ssim.update(predicted_label, label)
metrics = {'valid_loss': loss, 'valid_psnr': self.valid_psnr, 'valid_ssim': self.valid_ssim}
self.log_dict(metrics, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
loss, label, predicted_label = self._shared_step(batch)
self.test_psnr.update(predicted_label, label)
self.test_ssim.update(predicted_label, label)
metrics = {'test_loss': loss, 'test_psnr': self.test_psnr, 'test_ssim': self.test_ssim}
self.log_dict(metrics, on_epoch=True, on_step=False, prog_bar=True)
def configure_optimizers(self):
# We will support AdamW or SGD as optimizers.
if self.hparams.optimizer_name == "AdamW":
optimizer = torch.optim.AdamW(self.parameters(), **self.hparams.optimizer_hparams)
elif self.hparams.optimizer_name == "SGD":
optimizer = torch.optim.SGD(self.parameters(), **self.hparams.optimizer_hparams)
else:
assert False, f'Unknown optimizer: "{self.hparams.optimizer_name}"'
# scheduler = torch.optim.swa_utils.SWALR(optimizer, anneal_strategy="cos", anneal_epochs=5, swa_lr=0.002)
# scheduler = {'scheduler': torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95), 'name': 'expo_lr'}
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=10)
return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "valid_loss"}
